import re
from operator import itemgetter

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda

from AliUtils import get_model

# 创建模型
model = get_model()

# sqlalchemy 初始化MySQL数据库的连接
HOSTNAME = 'localhost'
PORT = '3306'
DATABASE = 'ry-vue'
USERNAME = 'root'
PASSWORD = 'root'
# mysqlclient驱动URL
# MYSQL_URI = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

# pymysql驱动URL
MYSQL_URI = 'mysql+pymysql://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)

# 测试连接是否成功
# print(db.get_usable_table_names())
# print(db.run('select * from sys_user limit 10;'))

# 直接使用大模型和数据库整合, 只能根据你的问题生成SQL
# 初始化生成SQL的chain
test_chain = create_sql_query_chain(model, db)
# resp = test_chain.invoke({'question': '请问：用户表中有多少条数据？'})
# print(resp)

answer_prompt = PromptTemplate.from_template(
	"""给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题。
	Question: {question}
	SQL Query: {query}
	SQL Result: {result}
	回答: """
)
# 创建一个执行sql语句的工具
execute_sql_tool = QuerySQLDataBaseTool(db=db)


def clean_query(query_str):
	# 使用正则表达式移除"SQLQuery: "前缀
	cleaned_query = re.sub(r'^SQLQuery:\s*', '', query_str).strip()
	return cleaned_query


# 1、生成SQL，2、执行SQL
# 2、模板
chain = (RunnablePassthrough
         .assign(query=test_chain)
         .assign(query_cleaned=RunnableLambda(lambda x: clean_query(x['query'])))
         .assign(result=itemgetter('query_cleaned') | execute_sql_tool)
         | answer_prompt
         | model
         | StrOutputParser()
         )

rep = chain.invoke(input={'question': '请问：员工表中有多少条数据？'})
print(rep)
